# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：visualize.py
    Created on：2025/2/21 17:37
    Author：(input)
    Description:势流研究中的可视化绘图代码
"""
import numpy as np
import matplotlib.pyplot as plt

# 设置全局字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示为方框的问题

# 网格绘图类
class GridPlotter:
    def __init__(self, grid):
        """
        初始化 GridPlotter 类，接收一个 Grid 对象
        """
        self.grid = grid

    def plot_grid(self):
        """
        绘制网格坐标并标注原点
        """
        plt.figure(figsize=(10, 10))

        # 绘制网格节点
        plt.scatter(self.grid.X, self.grid.Y, color='blue', marker='o', label="Grid Points")

        # 绘制网格线
        for i in range(self.grid.N):
            plt.plot(self.grid.X[i, :], self.grid.Y[i, :], color='gray', linestyle='--', alpha=0.5)  # 横向网格线
            plt.plot(self.grid.X[:, i], self.grid.Y[:, i], color='gray', linestyle='--', alpha=0.5)  # 纵向网格线

        # 标注原点
        origin_x, origin_y = self.grid.origin
        plt.scatter(0, 0, color='red', marker='x', s=100, label="Origin")

        # 显示原点坐标文本
        #plt.text(origin_x + 0.1, origin_y + 0.1, f"({origin_x}, {origin_y})", color='red', fontsize=12)

        # 设置标题和标签
        plt.title(f"网格Grid with {self.grid.N} x {self.grid.N} Points", fontsize=14)
        plt.xlabel('X Coordinate', fontsize=12)
        plt.ylabel('Y Coordinate', fontsize=12)

        # 显示网格
        plt.grid(True)
        plt.legend()
        plt.axis('equal')
        plt.show()

# 90°拐角流绘图类
class CornerFlowPlotter:
    """
    CornerFlowPlotter 用于绘制拐角流的物理现象，如势函数、流函数、速度场等。

    该类接收 Grid 提供的网格信息，以及 CornerFlow 提供的流动物理信息。
    """

    def __init__(self, grid, corner_flow):
        """
        初始化 CornerFlowPlotter

        参数：
        ----------
        grid: Grid
            网格信息
        corner_flow: CornerFlow
            提供拐角流的流动参数，如流动强度和源点位置
        """
        # 存储网格对象
        self.grid = grid
        # 获取网格坐标
        self.X = self.grid.X
        self.Y = self.grid.Y
        self.Z = self.grid.to_complex()

        # 计算绘图范围（放大 1.1 倍以适应可视化）
        self.x_range = (self.X.min(), self.X.max())
        self.y_range = (self.Y.min(), self.Y.max())
        self.plot_x_range = 1.1 * self.x_range[1]
        self.plot_y_range = 1.1 * self.y_range[1]
        self.figsize = (3*self.plot_x_range, 3*self.plot_y_range)

        # 存储拐角流对象
        self.corner_flow = corner_flow
        # 获取拐角坐标位置
        self.x_corner = corner_flow.x_cf
        self.y_corner = corner_flow.y_cf

    def plot_real_geo(self):
        """
        绘制拐角流的物理环境，包括源点和墙面边界。
        """
        # 拐角位置示意
        plt.scatter(self.x_corner, self.y_corner, color='r', s=80, marker="o", label='Corner Source')
        # 实体壁面位置示意
        plt.axhline(0., color='k', linestyle='-', linewidth=1)  # 下墙面
        plt.axvline(0., color='k', linestyle='-', linewidth=1)  # 左墙面
        #plt.legend()

    def plot_contour(self, data, levels, cmap, colorbar_label, title, xlabel='Width', ylabel='Height'):
        """
        通用等值线绘制方法，适用于绘制流函数 (psi)、势函数 (phi) 和其他场信息。

        参数：
        ----------
        data: np.ndarray
            需要绘制的二维数据（如势函数phi或流函数psi）
        levels: list
            等值线的级别
        cmap: str
            颜色映射（colormap）
        colorbar_label: str
            色条标签
        title: str
            图像标题
        xlabel: str
            x 轴标签
        ylabel: str
            y 轴标签
        """
        plt.figure(figsize=self.figsize)
        self.plot_real_geo()  # 绘制拐角流环境

        contour = plt.contour(self.X, self.Y, data, levels=levels, cmap=cmap, linewidths=2)
        plt.clabel(contour, inline=True, fontsize=10)
        plt.colorbar(contour, label=colorbar_label)

        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        plt.grid(True)
        plt.show()

    def plot_psi(self, psi, title='Stream Function'):
        """绘制流函数（ψ）的等值线图"""
        self.plot_contour(psi, levels=np.linspace(0.1, 4, 10),
                          cmap="coolwarm",
                          colorbar_label='Stream Function',
                          title=title)

    def plot_phi(self, phi, title='Potential Function'):
        """绘制势函数（φ）的等值线图"""
        self.plot_contour(phi, levels=np.linspace(-4, 4, 10),
                          cmap="viridis",
                          colorbar_label='Potential Function',
                          title=title)

    def plot_flow_net(self, phi, psi, title='Flow Net'):
        """
        绘制流网图，展示势函数和流函数的等值线。
        """
        plt.figure(figsize=self.figsize)

        contour_phi = plt.contour(self.X, self.Y, phi,
                                  levels=np.linspace(-4, 4, 10),
                                  cmap="viridis", linewidths=2)
        plt.clabel(contour_phi, inline=True, fontsize=10)

        contour_psi = plt.contour(self.X, self.Y, psi,
                                  levels=np.linspace(0.1, 4, 10),
                                  cmap="coolwarm", linewidths=2)
        plt.clabel(contour_psi, inline=True, fontsize=10)

        self.plot_real_geo()
        plt.xlabel('width (m)')
        plt.ylabel('height (m)')
        plt.title(title)
        plt.grid(True)
        plt.show()

    def plot_velocity_field(self, vx, vy, magV=None, show_magnitude=False, title='Velocity Field', cmap='jet'):
        """
        绘制速度场流线图，用箭头表示方向，可选通过颜色和颜色条表示速度大小。

        Parameters:
        ----------
        vx : np.ndarray
            水平速度分量
        vy : np.ndarray
            垂直速度分量
        magV : np.ndarray, optional
            速度模（大小），当 show_magnitude=True 时必须提供
        show_magnitude : bool, optional
            是否显示速度大小，默认 False
        title : str, optional
            图表标题，默认 'Velocity Field'
        cmap : str, optional
            颜色映射，默认 'jet'，仅在 show_magnitude=True 时生效
        """
        # 检查数据
        if not (np.issubdtype(vx.dtype, np.number) and np.issubdtype(vy.dtype, np.number)):
            raise ValueError("vx and vy must contain numeric values")
        if np.min(vx) == np.max(vx) and np.min(vy) == np.max(vy):
            raise ValueError("vx and vy must have a range of values for stream plotting")

        if show_magnitude:
            if magV is None:
                raise ValueError("magV must be provided when show_magnitude is True")
            if not np.issubdtype(magV.dtype, np.number):
                raise ValueError("magV must contain numeric values")
            if np.min(magV) == np.max(magV):
                raise ValueError("magV must have a range of values for color mapping")

        # 创建绘图
        plt.figure(figsize=self.figsize)
        if show_magnitude:
            stream = plt.streamplot(self.X, self.Y, vx, vy,
                                    density=1, linewidth=1,
                                    arrowsize=1, arrowstyle='->',
                                    color=magV, cmap=cmap)
            plt.colorbar(stream.lines, label='Velocity Magnitude (m/s)')
        else:
            stream = plt.streamplot(self.X, self.Y, vx, vy,
                                    density=1, linewidth=1,
                                    arrowsize=1, arrowstyle='->')

        self.plot_real_geo()
        plt.xlim(self.grid.X_ary[0], self.grid.X_ary[1])
        plt.ylim(self.grid.Y_ary[0], self.grid.Y_ary[1])
        plt.xlabel('X (m)')
        plt.ylabel('Y (m)')
        plt.title(title)
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.show()

    def plot_pressure(self, p, title='Pressure Distribution'):
        """
        绘制压力分布等值填充图。
        """
        plt.figure(figsize=self.figsize)
        contour = plt.contourf(self.X, self.Y, p,
                               levels=20, cmap='viridis')
        plt.colorbar(contour, label='Pressure', extend="both")

        self.plot_real_geo()
        plt.xlabel('Width (m)')
        plt.ylabel('Height (m)')
        plt.title(title)
        plt.grid(True)
        plt.show()

# 源汇流绘图类
class SosiFlowPlotter:
    """
    SosiFlowPlotter 用于绘制源汇流的物理现象，如流函数、势函数、速度场等。

    该类接收 Grid 提供的网格信息和 SourceSink 提供的流动参数。
    """

    def __init__(self, grid, source_sink):
        """
        初始化 SosiFlowPlotter

        Parameters:
        ----------
        grid : Grid
            提供网格信息的 Grid 对象
        source_sink : SourceSink
            包含源和汇流动参数的 SourceSink 对象
        """
        # 存储网格和源汇对象
        self.grid = grid
        self.source_sink = source_sink

        # 获取网格坐标
        self.X = self.grid.X
        self.Y = self.grid.Y
        self.Z = self.grid.to_complex()

        # 计算绘图范围（放大 1.1 倍）
        self.x_range = (self.X.min(), self.X.max())
        self.y_range = (self.Y.min(), self.Y.max())
        self.plot_x_range = 1.1 * self.x_range[1]
        self.plot_y_range = 1.1 * self.y_range[1]
        self.aspect_ratio = (self.grid.X_ary[1] - self.grid.X_ary[0]) / (self.grid.Y_ary[1] - self.grid.Y_ary[0])

        # 计算统一的 figsize
        base_width = 10  # 基础宽度（英寸），可调整
        base_height = base_width / self.aspect_ratio  # 根据宽高比计算高度
        self.figsize = (base_width, base_height)

        # 计算源和汇的坐标位置
        self.source_sink.compute_z_coordinates()
        self.source_x = np.real(self.source_sink.z_so)
        self.source_y = np.imag(self.source_sink.z_so)
        self.sink_x = np.real(self.source_sink.z_si)
        self.sink_y = np.imag(self.source_sink.z_si)

    def plot_real_geo(self):
        """
        绘制源汇流的物理环境，包括源点、汇点和房间边界。
        """
        # 源、汇位置示意
        plt.scatter(self.source_x, self.source_y, color='r', s=80, marker='o', label='Source')
        plt.scatter(self.sink_x, self.sink_y, color='g', s=80, marker='X', label='Sink')

        # 实体壁面位置示意
        wall_left = self.grid.X_ary[0]
        wall_right = self.grid.X_ary[1]
        wall_bottom = self.grid.Y_ary[0]
        wall_top = self.grid.Y_ary[1]
        plt.plot([wall_left, wall_right], [wall_top, wall_top], 'k-', lw=1)  # 上边界
        plt.plot([wall_left, wall_right], [wall_bottom, wall_bottom], 'k-', lw=1)  # 下边界
        plt.plot([wall_left, wall_left], [wall_bottom, wall_top], 'k-', lw=1)  # 左边界
        plt.plot([wall_right, wall_right], [wall_bottom, wall_top], 'k-', lw=1)  # 右边界
        plt.legend(loc='lower right')

    def plot_contour(self, data, levels, cmap, colorbar_label, title, xlabel='Width', ylabel='Height', num_levels=12):
        """
        通用等值线绘制方法，适用于绘制流函数 (psi)、势函数 (phi) 和其他场信息。
        self, data, title, cmap, colorbar_label, levels=None, num_levels=12
        Parameters:
        ----------
        data: np.ndarray
            需要绘制的二维数据（如势函数phi或流函数psi）
        levels: list
            等值线的级别
        cmap: str
            颜色映射（colormap）
        colorbar_label: str
            色条标签
        title: str
            图像标题
        xlabel: str
            x 轴标签
        ylabel: str
            y 轴标签
        """
        # 检查数据类型
        if not np.issubdtype(data.dtype, np.number):
            raise ValueError(f"{colorbar_label} must contain numeric values")
        min_val, max_val = np.min(data), np.max(data)
        if min_val == max_val:
            raise ValueError(f"{colorbar_label} must have a range of values for contouring")

        plt.figure(figsize=self.figsize)
        self.plot_real_geo()# 绘制源汇位置与墙壁位置

        levels = levels if levels is not None else np.linspace(min_val, max_val, num_levels)
        contour = plt.contour(self.X, self.Y, data, levels=levels,
                              cmap=cmap, linewidths=2)
        plt.clabel(contour, inline=True, fontsize=10, fmt="%.2e")
        plt.colorbar(contour, label=colorbar_label)

        self.contour_set = contour# 这允许后续方法（如 extract_contour_coordinates）访问等值线数据。

        plt.xlim(self.grid.X_ary[0], self.grid.X_ary[1])
        plt.ylim(self.grid.Y_ary[0], self.grid.Y_ary[1])

        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        plt.grid(True)
        plt.show()

    def extract_contour_coordinates(self, contour_value):
        """
        从等值线中提取指定值的坐标点。

        Parameters:
        ----------
        contour_value : float
            要提取的等值线水平值

        Returns:
        -------
        coords : list
            包含该等值线所有坐标点的列表
        """
        if not hasattr(self, 'contour_set'):
            raise AttributeError("No contour set available. Run a contour plot method first.")
        coords = []
        for path in self.contour_set.get_paths():
            vertices = path.vertices
            if np.isclose(vertices[:, 0], contour_value).any() or np.isclose(vertices[:, 1], contour_value).any():
                coords.append(vertices)
        return coords

    def plot_psi(self, psi, title='Stream Function'):
        """绘制流函数的等值线图"""
        #data, levels, cmap, colorbar_label, title, xlabel = 'Width', ylabel = 'Height', num_levels = 12
        self.plot_contour(psi, levels=None,
                          cmap="coolwarm",
                          colorbar_label='Stream Function',
                          title=title)

    def plot_phi(self, phi, title='Potential Function'):
        """绘制势函数的等值线图"""
        self.plot_contour(phi, levels=None,
                          cmap="coolwarm",
                          colorbar_label='Potential Function',
                          title=title)

    def plot_flow_net(self, phi, psi, title='Flow Net', num_levels=12):
        """
        绘制流网图，展示势函数和流函数的等值线。
        """
        # 检查数据
        if not (np.issubdtype(phi.dtype, np.number) and np.issubdtype(psi.dtype, np.number)):
            raise ValueError("phi and psi must contain numeric values")
        min_phi, max_phi = np.min(phi), np.max(phi)
        min_psi, max_psi = np.min(psi), np.max(psi)
        if min_phi == max_phi or min_psi == max_psi:
            raise ValueError("phi and psi must have a range of values for contouring")

        plt.figure(figsize=self.figsize)

        levels_phi = np.linspace(min_phi, max_phi, num_levels)
        levels_psi = np.linspace(min_psi, max_psi, num_levels)

        contour_phi = plt.contour(self.X, self.Y, phi,
                                  levels=levels_phi, cmap="viridis", linewidths=2)
        plt.clabel(contour_phi, inline=True, fontsize=10, fmt="%.2e")
        contour_psi = plt.contour(self.X, self.Y, psi,
                                  levels=levels_psi, cmap="coolwarm", linewidths=2)
        plt.clabel(contour_psi, inline=True, fontsize=10, fmt="%.2e")
        #self.contour_set = contour_psi  # 存储最后一个等值线集合

        self.plot_real_geo()

        plt.xlim(self.grid.X_ary[0], self.grid.X_ary[1])
        plt.ylim(self.grid.Y_ary[0], self.grid.Y_ary[1])
        plt.xlabel(xlabel='width')
        plt.ylabel(ylabel='height')
        plt.title(title)
        plt.grid(True)
        plt.show()

    def plot_velocity_field(self, vx, vy, magV=None, show_magnitude=False, title='Velocity Field', cmap='jet'):
        """
        绘制速度场流线图，用箭头表示方向，可选通过颜色和颜色条表示速度大小。

        Parameters:
        ----------
        vx : np.ndarray
            水平速度分量
        vy : np.ndarray
            垂直速度分量
        magV : np.ndarray, optional
            速度模（大小），当 show_magnitude=True 时必须提供
        show_magnitude : bool, optional
            是否显示速度大小，默认 False
        title : str, optional
            图表标题，默认 'Velocity Field'
        cmap : str, optional
            颜色映射，默认 'jet'，仅在 show_magnitude=True 时生效
        """
        # 检查数据
        if not (np.issubdtype(vx.dtype, np.number) and np.issubdtype(vy.dtype, np.number)):
            raise ValueError("vx and vy must contain numeric values")
        if np.min(vx) == np.max(vx) and np.min(vy) == np.max(vy):
            raise ValueError("vx and vy must have a range of values for stream plotting")

        if show_magnitude:
            if magV is None:
                raise ValueError("magV must be provided when show_magnitude is True")
            if not np.issubdtype(magV.dtype, np.number):
                raise ValueError("magV must contain numeric values")
            if np.min(magV) == np.max(magV):
                raise ValueError("magV must have a range of values for color mapping")

        # 创建绘图
        plt.figure(figsize=self.figsize)
        if show_magnitude:
            stream = plt.streamplot(self.X, self.Y, vx, vy,
                                    density=1, linewidth=1,
                                    arrowsize=1, arrowstyle='->',
                                    color=magV, cmap=cmap)
            plt.colorbar(stream.lines, label='Velocity Magnitude (m/s)')
        else:
            stream = plt.streamplot(self.X, self.Y, vx, vy,
                                    density=1, linewidth=1,
                                    arrowsize=1, arrowstyle='->')

        self.plot_real_geo()
        plt.xlim(self.grid.X_ary[0], self.grid.X_ary[1])
        plt.ylim(self.grid.Y_ary[0], self.grid.Y_ary[1])
        plt.xlabel('width')
        plt.ylabel('height')
        plt.title(title)
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()  # 自动调整布局
        plt.show()


    def plot_pressure(self, p, title='Pressure Coefficient Distribution', num_levels=20):
        """
        绘制压力分布等值线图。
        """
        # 检查数据
        if not np.issubdtype(p.dtype, np.number):
            raise ValueError("Pressure must contain numeric values")
        min_p, max_p = np.min(p), np.max(p)
        if min_p == max_p:
            raise ValueError("Pressure must have a range of values for contouring")

        plt.figure(figsize=self.figsize)
        contour = plt.contour(self.X, self.Y, p,
                              levels=num_levels, cmap='viridis', linewidths=1)
        self.contour_set = contour
        plt.clabel(contour, inline=True, fontsize=10, fmt="%.2e")
        self.plot_real_geo()

        plt.colorbar(contour, label='Pressure Coefficient (Cp)', extend="both")
        plt.xlim(self.grid.X_ary[0], self.grid.X_ary[1])
        plt.ylim(self.grid.Y_ary[0], self.grid.Y_ary[1])
        plt.xlabel('Width')
        plt.ylabel('Height')
        plt.title(title)
        plt.grid(True)
        plt.show()

"""____示例____"""
# Example of using the general models
if __name__ == "__main__":
    print("finish")
